# 根据对应行内容搜索链接代码，随机个数（一个excel找3-4个链接，可调频）

import requests
import os
from typing import Dict, Any, Optional
import openai

import openai
from typing import List


# API设置
OPENAI_API_KEY = "your-api-key"  # 替换为你的OpenAI API密钥
OPENAI_API_BASE = "https://www.dmxapi.com/v1"  # 可以替换为其他API端点
SERPER_API_KEY = "your-serper-key"  # 替换为你的Serper API密钥
SERPER_API_URL = "https://google.serper.dev/search"
OPENAI_MODEL = "gpt-3.5-turbo"  # 可以根据需要更换模型

class EnhancedLinkFinder:
    def __init__(self, serper_api_key: str, openai_api_key: str, openai_api_base: str = OPENAI_API_BASE):
        """初始化增强版链接查找器，集成OpenAI生成搜索关键词"""
        self.serper_api_key = serper_api_key
        self.openai_api_key = openai_api_key
        self.openai_api_base = openai_api_base
        openai.api_base = self.openai_api_base  # 设置自定义API端点
        openai.api_key = openai_api_key  # 设置API密钥
        
    def search_google(self, query: str) -> Dict[str, Any]:
        """使用Serper API搜索Google"""
        headers = {
            'X-API-KEY': self.serper_api_key,
            'Content-Type': 'application/json'
        }
        
        payload = {
            'q': query,
            'num': 3  # 请求3个结果以增加找到相关信息的可能性
        }
        
        try:
            response = requests.post(
                SERPER_API_URL, 
                headers=headers,
                json=payload
            )
            response.raise_for_status()  # 检查HTTP错误
            return response.json()
        except Exception as e:
            print(f"搜索错误: {e}")
            return {"error": str(e)}

    def generate_search_keywords(
        self,
        text: str,
        openai_api_key: str = 'sk-taGSEBrnb1G0maRcE7ZenNFnjrHBNgAiHSsfr8FVTAbZpkUH',
        openai_api_base: str = "https://www.dmxapi.com/v1",
        model: str = "gpt-3.5-turbo",
        num_keywords: int = 5,
        max_tokens: int = 100,
        temperature: float = 0.3
    ) -> List[str]:
        """
        从给定文本中提取最重要的关键词，用于搜索引擎查询。

        参数：
            text: 要分析的原始文本。
            openai_api_key: OpenAI API Key（形如 sk-xxxx）。
            openai_api_base: API 服务端点，默认为官方 https://api.openai.com/v1。
            model: 使用的模型名称。
            num_keywords: 希望提取的关键词个数。
            max_tokens: 接口返回最大 token 数。
            temperature: 生成的随机性温度。

        返回：
            一个关键词列表，已经去除多余空格。
        """

        # 全局配置
        openai.api_key  = openai_api_key
        openai.api_base = openai_api_base

        # 构造系统和用户消息
        system_prompt = (
            f"你是一个搜索优化助手。"
            f"请从以下文本中提取 {num_keywords} 个最重要的关键词或短语，用于搜索引擎查询。"
            "关键词应该简洁、相关，并能代表文本核心内容。"
            "输出格式为逗号分隔的关键词列表，不要添加其他解释。"
        )
        messages = [
            {"role": "system",  "content": system_prompt},
            {"role": "user",    "content": text}
        ]

        # 调用 OpenAI ChatCompletion
        response = openai.ChatCompletion.create(
            model= model,
            messages= messages,
            max_tokens= max_tokens,
            temperature= temperature
        )

        # 解析并清洗结果
        raw = response.choices[0].message.content.strip()
        keywords = [kw.strip() for kw in raw.split(",") if kw.strip()]

        return keywords

    
    def extract_best_link(self, search_results: Dict[str, Any]) -> Optional[Dict[str, str]]:
        """从搜索结果中提取最佳链接"""
        # 首先检查organic结果
        if "organic" in search_results and search_results["organic"]:
            item = search_results["organic"][0]  # 获取第一个结果
            if "title" in item and "link" in item:
                return {
                    "title": item["title"],
                    "link": item["link"],
                    "snippet": item.get("snippet", "")
                }
        
        # 如果没有organic结果，检查知识图谱
        if "knowledgeGraph" in search_results:
            kg = search_results["knowledgeGraph"]
            if "title" in kg and "url" in kg:
                return {
                    "title": f"知识图谱: {kg['title']}",
                    "link": kg["url"],
                    "snippet": kg.get("description", "")
                }
        
        return None  # 没有找到链接
    
    def find_link(self, text: str) -> Optional[Dict[str, str]]:
        """为文本查找一个相关链接"""
        if len(text) > 50:
            # 如果文本较长，使用OpenAI生成搜索关键词
            try:
                keywords = self.generate_search_keywords(text)
                search_query = ", ".join(keywords)  # 将关键词列表拼接成搜索查询字符串
                print(f"生成的搜索关键词: {search_query}")
            except Exception as e:
                print(f"生成搜索关键词时出错: {e}")
                return None
        
            # 使用生成的关键词搜索
            search_results = self.search_google(search_query)
        else:
            # 如果文本较短，直接使用原文本进行搜索
            search_results = self.search_google(text)
        
        # 提取最佳链接
        best_link = self.extract_best_link(search_results)
        return best_link


import os
import openpyxl
from typing import Optional
import random

def process_xlsx_files_recursively(folder_path: str, finder) -> None:
    """
    遍历文件夹下所有xlsx文件，读取第3行第5列内容，调用finder.find_link()
    参数：
        folder_path: 根目录路径
        finder: 已初始化的 EnhancedLinkFinder 实例
    """
    for root, _, files in os.walk(folder_path):  
        for file in files:
            if file.endswith(".xlsx"):
                file_path = os.path.join(root, file)
                try:
                    workbook = openpyxl.load_workbook(file_path, data_only=True)
                    sheet = workbook.active  # 默认处理第一个sheet
                    
                    # 获取最大行数，用于遍历整个表格
                    max_row = sheet.max_row
                    print(f"\n📄 正在处理文件: {file_path}")
                    print(f"总行数: {max_row}")
                    
                    # 从第1行开始
                    row = 1
                    while row <= max_row:
                        # 获取第1列和第5列的值
                        col1_value = sheet.cell(row=row, column=1).value
                        col5_value = sheet.cell(row=row, column=5).value
                        
                        # 检查第5列是否有有效内容
                        if col5_value and isinstance(col5_value, str) and col5_value.strip():
                            print(f"行 {row}: 第5列内容: {col5_value.strip()}")
                            
                            # 检查第1列的值是否允许填写链接
                            if col1_value and is_link_allowed(col1_value):
                                # 如果第5列内容为"NULL"，则跳过链接查找
                                if col5_value.strip() == "NULL":
                                    print(f"⚠️ 行 {row}: 第5列内容为NULL，跳过链接查找")
                                else:
                                    # 查找链接
                                    link = finder.find_link(col5_value.strip())
                                    if link:
                                        print(f"✅ 行 {row}: 找到相关链接")
                                        # 将链接填写到第6列
                                        sheet.cell(row=row, column=6).value = link['link']
                                        print(f"已将链接 {link['link']} 填写到第6列")
                                    else:
                                        print(f"⚠️ 行 {row}: 未找到相关链接")
                            else:
                                print(f"⚠️ 行 {row}: 根据第1列的值 '{col1_value}'，不允许填写链接")
                        
                        # 生成随机步长，决定下次处理哪一行
                        random_step = random.randint(1, 7)
                        print(f"下一次将跳过 {random_step} 行")
                        row += random_step
                    
                    # 保存更改到Excel文件
                    workbook.save(file_path)
                    print(f"✅ 文件 {file_path} 处理完成并已保存")
                    
                except Exception as e:
                    print(f"❌ 处理文件 {file_path} 时出错: {e}")


def is_link_allowed(col1_value):
    """
    根据第1列的值判断是否允许填写链接
    
    Args:
        col1_value: 第1列的值
        
    Returns:
        bool: 如果允许填写链接返回True，否则返回False
    """
    # 定义不允许生成链接的列表
    staticListCODE = ['分类码', '图片']
    staticListNULL = ['区域代码', '国家或行业标准编码', '军队内部编码', '军用基础地理信息编码', 
                      '地理标示', '关联地图图种及比例尺', '外挂表名称', '视频', '地图', 
                      '关联地图图种及比例尺', '视频数据','数据项','顺序码','名称']
    
    # 合并两个列表
    restricted_values = staticListCODE + staticListNULL
    
    # 检查第1列的值是否在不允许生成链接的列表中
    if col1_value and isinstance(col1_value, str):
        # 去除可能的空格并转换为小写进行比较，以增强匹配的准确性
        col1_value_cleaned = col1_value.strip().lower()
        
        for restricted_value in restricted_values:
            if col1_value_cleaned == restricted_value.lower():
                return False
    
    # 如果不在限制列表中，则允许生成链接
    return True


# 使用示例
def main():
    # 配置OpenAI
    openai.api_key = OPENAI_API_KEY  # 设置API密钥
    openai.api_base = OPENAI_API_BASE  # 设置自定义API端点

    # 使用环境变量或直接设置API密钥
    serper_api_key = os.environ.get("SERPER_API_KEY", SERPER_API_KEY)
    openai_api_key = os.environ.get("OPENAI_API_KEY", OPENAI_API_KEY)
    

    # 创建增强版链接查找器
    finder = EnhancedLinkFinder(serper_api_key, openai_api_key)

    # 读取文件夹下所有xlsx并处理
    folder_path = r"C:\Users\xsr\Desktop\me\加勒比海数据_本组\基底数据"  # ← 请替换为你的文件夹路径
    process_xlsx_files_recursively(folder_path, finder)

    # # 获取用户输入
    # print("请输入要查找链接的文本 (可以是较长的文本):")
    # text = input()

    # # 查找相关链接
    # link = finder.find_link(text)
    
    # # 打印结果
    # if link:
    #     print("\n找到的相关链接:")
    #     print(f"标题: {link['title']}")
    #     print(f"URL: {link['link']}")
    #     print(f"摘要: {link.get('snippet', '')}")
    # else:
    #     print("\n未找到相关链接")

if __name__ == "__main__":
    main()